from flask import Flask, render_template, request, jsonify, Response
from flask_cors import CORS, cross_origin
import time
import logging
import cv2
from djitellopy import tello

# 初始化Flask应用
app = Flask(__name__)

# 配置CORS - 允许所有来源的跨域请求
# 生产环境中建议限制具体的前端域名
CORS(app, resources={
    r"/*": {
        "origins": ["*"],  # 允许所有来源
        "methods": ["GET", "POST", "OPTIONS"],  # 允许的HTTP方法
        "allow_headers": ["Content-Type", "Authorization"]  # 允许的请求头
    }
})

# 全局变量
drone = None
is_drone_connected = False
drone_speed = 70
font = cv2.FONT_HERSHEY_SIMPLEX
fontScale = 0.5
lineThickness = 1

# 录像相关全局变量
is_recording = False
video_writer = None
output_video = None


# 初始化无人机连接
def init_drone():
    global drone, is_drone_connected
    try:
        drone = tello.Tello()
        drone.connect()
        print(f"无人机电池电量: {drone.get_battery()}%")
        drone.streamon()
        drone.LOGGER.setLevel(logging.ERROR)
        time.sleep(2)  # 等待视频初始化
        is_drone_connected = True
        print("无人机连接成功!")
        return True
    except Exception as e:
        print(f"无人机连接失败: {e}")
        return False


# 生成视频流的函数
def generate_frames():
    global drone, is_drone_connected, is_recording, video_writer

    while is_drone_connected:
        try:
            if drone and drone.stream_on:
                # 从无人机获取原始帧 (BGR格式)
                img = drone.get_frame_read().frame
                bgr_frame = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

                if bgr_frame is not None:
                    # 添加状态信息到视频帧
                    battery = drone.get_battery() if drone else 0
                    height = drone.get_height() if drone else 0
                    info_text = f"Battery: {battery}% Height: {height}cm Time: {time.strftime('%H:%M:%S')}"

                    # 如果正在录像，添加录像标识
                    if is_recording:
                        info_text += " | RECORDING"

                    # 在BGR帧上添加文字
                    cv2.putText(bgr_frame, info_text, (20, 30), font, fontScale, (0, 0, 255), lineThickness)

                    # 转换为JPEG格式用于流式传输
                    ret, buffer = cv2.imencode('.jpg', bgr_frame)

                    # 如果正在录像，将BGR帧写入视频文件
                    if is_recording and video_writer:
                        video_writer.write(bgr_frame)

                    if ret:
                        frame_bytes = buffer.tobytes()
                        yield (b'--frame\r\n'
                               b'Content-Type: image/jpeg\r\n\r\n' + frame_bytes + b'\r\n')

            time.sleep(0.03)  # 约30FPS
        except Exception as e:
            print(f"视频流错误: {e}")
            time.sleep(1)


# Flask路由
@app.route('/')
def index():
    return render_template('index.html')


@app.route('/vue')
def vue_control():
    return render_template('vue-control.html')


@app.route('/video_feed')
def video_feed():
    return Response(generate_frames(),
                    mimetype='multipart/x-mixed-replace; boundary=frame')


@app.route('/connect_drone', methods=['POST', 'OPTIONS'])
@cross_origin()
def connect_drone():
    if request.method == 'OPTIONS':
        return jsonify({'status': 'success'})

    global is_drone_connected
    if init_drone():
        is_drone_connected = True
        return jsonify({'status': 'success', 'message': '无人机连接成功'})
    else:
        return jsonify({'status': 'error', 'message': '无人机连接失败'})


@app.route('/disconnect_drone', methods=['POST', 'OPTIONS'])
@cross_origin()
def disconnect_drone():
    if request.method == 'OPTIONS':
        return jsonify({'status': 'success'})

    global drone, is_drone_connected, is_recording, video_writer

    try:
        # 如果正在录像，先停止录像
        if is_recording and video_writer:
            video_writer.release()
            video_writer = None
            is_recording = False

        if drone:
            if is_drone_connected:
                drone.land()
                time.sleep(2)
            drone.streamoff()
            drone.end()

        is_drone_connected = False
        return jsonify({'status': 'success', 'message': '无人机已断开连接'})
    except Exception as e:
        return jsonify({'status': 'error', 'message': f'断开连接时出错: {e}'})


@app.route('/control', methods=['POST', 'OPTIONS'])
@cross_origin()
def control_drone():
    if request.method == 'OPTIONS':
        return jsonify({'status': 'success'})

    global drone, is_drone_connected, drone_speed

    if not is_drone_connected:
        return jsonify({'status': 'error', 'message': '无人机未连接'})

    data = request.json
    command = data.get('command')

    try:
        if command == 'takeoff':
            drone.takeoff()
            return jsonify({'status': 'success', 'message': '起飞命令已发送'})

        elif command == 'land':
            drone.land()
            return jsonify({'status': 'success', 'message': '降落命令已发送'})

        elif command == 'emergency':
            drone.emergency()
            return jsonify({'status': 'success', 'message': '紧急停止命令已发送'})

        elif command == 'move':
            direction = data.get('direction')
            lr, fb, ud, yv = 0, 0, 0, 0

            if direction == 'left':
                lr = -drone_speed
            elif direction == 'right':
                lr = drone_speed
            elif direction == 'forward':
                fb = drone_speed
            elif direction == 'backward':
                fb = -drone_speed
            elif direction == 'up':
                ud = drone_speed
            elif direction == 'down':
                ud = -drone_speed
            elif direction == 'rotate_left':
                yv = -drone_speed
            elif direction == 'rotate_right':
                yv = drone_speed
            elif direction == 'stop':
                lr, fb, ud, yv = 0, 0, 0, 0

            drone.send_rc_control(lr, fb, ud, yv)
            return jsonify({'status': 'success', 'message': f'移动命令: {direction}'})

        return jsonify({'status': 'error', 'message': '未知命令'})

    except Exception as e:
        return jsonify({'status': 'error', 'message': str(e)})


@app.route('/take_photo', methods=['POST', 'OPTIONS'])
@cross_origin()
def take_photo():
    if request.method == 'OPTIONS':
        return jsonify({'status': 'success'})

    global drone, is_drone_connected
    if not is_drone_connected:
        return jsonify({'status': 'error', 'message': '无人机未连接'})

    try:
        # 获取原始BGR帧
        img_photo = drone.get_frame_read().frame
        photo = cv2.cvtColor(img_photo, cv2.COLOR_BGR2RGB)

        if photo is not None:
            filename = f'snap-{time.strftime("%Y%m%d-%H%M%S")}.jpg'
            # 保存BGR格式的帧
            cv2.imwrite(filename, photo)
            return jsonify({'status': 'success', 'message': f'照片已保存为 {filename}'})
        return jsonify({'status': 'error', 'message': '获取画面失败'})
    except Exception as e:
        return jsonify({'status': 'error', 'message': f'拍照失败: {e}'})


@app.route('/start_recording', methods=['POST', 'OPTIONS'])
@cross_origin()
def start_recording():
    if request.method == 'OPTIONS':
        return jsonify({'status': 'success'})

    global is_recording, video_writer, output_video, drone, is_drone_connected

    if not is_drone_connected:
        return jsonify({'status': 'error', 'message': '无人机未连接'})

    if is_recording:
        return jsonify({'status': 'error', 'message': '已经在录像中'})

    try:
        # 创建视频文件名
        timestamp = time.strftime("%Y%m%d-%H%M%S")
        output_video = f'record-{timestamp}.mp4'

        # 获取视频帧信息以设置编码器
        bgr_frame = drone.get_frame_read().frame
        height, width, _ = bgr_frame.shape

        # 定义编码器并创建VideoWriter对象
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        video_writer = cv2.VideoWriter(output_video, fourcc, 30.0, (width, height))

        is_recording = True
        return jsonify({'status': 'success', 'message': f'开始录像，将保存为 {output_video}'})
    except Exception as e:
        return jsonify({'status': 'error', 'message': f'启动录像失败: {e}'})


@app.route('/stop_recording', methods=['POST', 'OPTIONS'])
@cross_origin()
def stop_recording():
    if request.method == 'OPTIONS':
        return jsonify({'status': 'success'})

    global is_recording, video_writer, output_video

    if not is_recording:
        return jsonify({'status': 'error', 'message': '没有在录像中'})

    try:
        is_recording = False
        if video_writer:
            video_writer.release()
            video_writer = None

        return jsonify({'status': 'success', 'message': f'录像已保存为 {output_video}'})
    except Exception as e:
        return jsonify({'status': 'error', 'message': f'停止录像失败: {e}'})


@app.route('/battery', methods=['GET', 'OPTIONS'])
@cross_origin()
def get_battery():
    if request.method == 'OPTIONS':
        return jsonify({'status': 'success'})

    global drone, is_drone_connected
    if not is_drone_connected:
        return jsonify({'status': 'error', 'message': '无人机未连接', 'battery': 0})

    try:
        battery = drone.get_battery()
        return jsonify({'status': 'success', 'battery': battery})
    except Exception as e:
        return jsonify({'status': 'error', 'message': str(e), 'battery': 0})


if __name__ == '__main__':
    print("启动Flask服务器...")
    print("请打开浏览器访问: http://localhost:5000")
    app.run(host='0.0.0.0', port=5000, debug=True, threaded=True)